package com.nxw.springcloud.gateway.config;

import com.alibaba.fastjson.JSON;
import com.nxw.springcloud.gateway.RedisService;
import com.nxw.springcloud.gateway.entity.EnvironmentEnum;
import com.nxw.springcloud.gateway.entity.RateLimitKey;
import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import org.springframework.beans.BeansException;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.cloud.gateway.filter.ratelimit.AbstractRateLimiter;
import org.springframework.context.ApplicationContext;
import org.springframework.context.ApplicationContextAware;
import org.springframework.data.redis.core.ReactiveRedisTemplate;
import org.springframework.data.redis.core.script.RedisScript;
import org.springframework.validation.Validator;
import org.springframework.validation.annotation.Validated;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;

import javax.validation.constraints.Min;
import javax.validation.constraints.NotNull;
import java.time.Instant;
import java.util.*;
import java.util.concurrent.atomic.AtomicBoolean;

public class RateRedisRateLimiter extends AbstractRateLimiter<RateRedisRateLimiter.Config> implements ApplicationContextAware {
 
    public static final String CONFIGURATION_PROPERTY_NAME = "redis-rate-limiter";
    public static final String REDIS_SCRIPT_NAME = "redisRequestRateLimiterScript";
    public static final String REMAINING_HEADER = "X-RateLimit-Remaining";
    public static final String REPLENISH_RATE_HEADER = "X-RateLimit-Replenish-Rate";
    public static final String BURST_CAPACITY_HEADER = "X-RateLimit-Burst-Capacity";
    private Log log = LogFactory.getLog(this.getClass());
    private ReactiveRedisTemplate<String, String> redisTemplate;
    private RedisScript<List<Long>> script;
    private AtomicBoolean initialized = new AtomicBoolean(false);
    private RateRedisRateLimiter.Config defaultConfig;
    private boolean includeHeaders = true;
    private String remainingHeader = "X-RateLimit-Remaining";
    private String replenishRateHeader = "X-RateLimit-Replenish-Rate";
    private String burstCapacityHeader = "X-RateLimit-Burst-Capacity";
 
    public volatile static Map<String, RateLimitConfig> rateLimitConfigMap = new HashMap<>();
 
    @Autowired
    RedisService redisConfigService;
 
    public RateRedisRateLimiter(ReactiveRedisTemplate<String, String> redisTemplate, RedisScript<List<Long>> script, Validator validator) {
        super(RateRedisRateLimiter.Config.class, "redis-rate-limiter", validator);
        this.redisTemplate = redisTemplate;
        this.script = script;
        this.initialized.compareAndSet(false, true);
    }
 
    public RateRedisRateLimiter(int defaultReplenishRate, int defaultBurstCapacity) {
        super(RateRedisRateLimiter.Config.class, "redis-rate-limiter", (Validator)null);
        this.defaultConfig = (new RateRedisRateLimiter.Config()).setReplenishRate(defaultReplenishRate).setBurstCapacity(defaultBurstCapacity);
    }
 
    public boolean isIncludeHeaders() {
        return this.includeHeaders;
    }
 
    public void setIncludeHeaders(boolean includeHeaders) {
        this.includeHeaders = includeHeaders;
    }
 
    public String getRemainingHeader() {
        return this.remainingHeader;
    }
 
    public void setRemainingHeader(String remainingHeader) {
        this.remainingHeader = remainingHeader;
    }
 
    public String getReplenishRateHeader() {
        return this.replenishRateHeader;
    }
 
    public void setReplenishRateHeader(String replenishRateHeader) {
        this.replenishRateHeader = replenishRateHeader;
    }
 
    public String getBurstCapacityHeader() {
        return this.burstCapacityHeader;
    }
 
    public void setBurstCapacityHeader(String burstCapacityHeader) {
        this.burstCapacityHeader = burstCapacityHeader;
    }
 
    public void setApplicationContext(ApplicationContext context) throws BeansException {
        if (this.initialized.compareAndSet(false, true)) {
            this.redisTemplate = (ReactiveRedisTemplate)context.getBean("stringReactiveRedisTemplate", ReactiveRedisTemplate.class);
            this.script = (RedisScript)context.getBean("redisRequestRateLimiterScript", RedisScript.class);
            if (context.getBeanNamesForType(Validator.class).length > 0) {
                this.setValidator((Validator)context.getBean(Validator.class));
            }
        }
 
    }
 
    RateRedisRateLimiter.Config getDefaultConfig() {
        return this.defaultConfig;
    }
 
    public Mono<Response> isAllowed(String routeId, String id) {
        if (!this.initialized.get()) {
            throw new IllegalStateException("RateRedisRateLimiter is not initialized");
        } else {
            RateLimitConfig limitConfig = getLimitConfig(routeId);
 
            RateLimitKey rateLimitKey = JSON.parseObject(id, RateLimitKey.class);
            if (limitConfig == null || limitConfig.getTokenConfig().size() == 0) {
                return Mono.just(new Response(true,null));
            }
 
            Map<String, Config> conf = limitConfig.getTokenConfig();
            // IP限流
            String clientIp = rateLimitKey.getClientIp();
            Config clientIpConfig = conf.get("clientIp");
            // APPkEY 限流
            String appKeyAndSecret = rateLimitKey.getAppKey();
            Config appKeyConfig = conf.get("appKeyAndSecret");
            Config appKetLimitConfig = redisConfigService.getAppKetLimit(appKeyAndSecret,appKeyConfig);

 
            // 校验IP访问过滤
            return isSingleAllow(clientIp, routeId, clientIpConfig).flatMap(
                ipRes -> {
                    if (!ipRes.isAllowed()) {
                        Map<String, String> header = new HashMap<>(ipRes.getHeaders());
                        header.put("limitReason", "您当前IP(" + clientIp + ")访问过快");
                        System.out.println("isAllow:" + ipRes.isAllowed());
                        System.out.println("您当前IP(" + clientIp + ")访问过快");
                        return Mono.just(new Response(ipRes.isAllowed(), header));
                    } else {
                        //appKey 限流
                        return isSingleAllow(appKeyAndSecret, routeId, appKetLimitConfig).flatMap(
                            serviceRes -> {
                                if (!serviceRes.isAllowed()) {
                                    Map<String, String> header = new HashMap<>();
                                    header.putAll(serviceRes.getHeaders());
                                    header.put("limitReason", "此appKey当前访问量过快，请稍候再试");
                                    System.out.println("isAllow:" + serviceRes.isAllowed());
                                    System.out.println("此appKey当前访问量过快，请稍候再试");
                                    return Mono.just(new Response(serviceRes.isAllowed(), header));
                                } else {
                                    // TODO 后续再加其他校验
                                    return Mono.just(serviceRes);
                                }
                            }
                        );
                    }
                }
            );
        }
    }
 
    /**
     * 单级检测限流
     * @param key
     * @param routeId
     * @param config
     * @return
     */
    private Mono<Response> isSingleAllow(String key, String routeId, Config config) {
        if (config == null) {
            return Mono.just(new Response(true, new HashMap<>()));
        }
        // How many requests per second do you want a user to be allowed to do?
        // 允许用户每秒处理多少个请求
        int replenishRate = config.getReplenishRate();
 
        // How much bursting do you want to allow?
        // 令牌桶的容量，允许在一秒钟内完成的最大请求数
        int burstCapacity = config.getBurstCapacity();
 
        try {
            List<String> keys = getKeys(routeId + "$" + key);
 
            // The arguments to the LUA script. time() returns unixtime in seconds.
            List<String> scriptArgs = Arrays.asList(replenishRate + "", burstCapacity + "",
                    getTime(config.timeType) + "", "1");
            // allowed, tokens_left = redis.eval(SCRIPT, keys, args)
            Flux<List<Long>> flux = this.redisTemplate.execute(this.script, keys, scriptArgs);
            // .log("redisratelimiter", Level.FINER);
            return flux.onErrorResume(throwable -> Flux.just(Arrays.asList(1L, -1L)))
                    .reduce(new ArrayList<Long>(), (longs, l) -> {
                        longs.addAll(l);
                        return longs;
                    }) .map(results -> {
                        boolean allowed = results.get(0) == 1L;
                        Long tokensLeft = results.get(1);
 
                        Response response = new Response(allowed, getHeaders(config, tokensLeft));
 
                        if (log.isDebugEnabled()) {
                            log.debug("response: " + response);
                        }
                        return response;
                    });
        }
        catch (Exception e) {
            /*
             * We don't want a hard dependency on Redis to allow traffic. Make sure to set
             * an alert so you know if this is happening too much. Stripe's observed
             * failure rate is 0.01%.
             */
            log.error("Error determining if user allowed from redis", e);
        }
        return Mono.just(new Response(true, getHeaders(config, -1L)));
    }


    public static  long getTime(int type){
        long time = Instant.now().getEpochSecond();
        switch (type){
            case 1:
                break;
            case 2:
                time = time / (60);
                break;
            case 3:
                time = time / (60 * 60);
                break;
            case 4:
                time = time / (60 * 60 * 24);
                break;
        }
        return time;
    }


    /**
     * 从redis中获取配置信息
     * @param routeId
     * @return
     */
    private RateLimitConfig getLimitConfig(String routeId) {
        if (!rateLimitConfigMap.containsKey(routeId)) {
            redisConfigService.refreshRateLimitConfig(routeId);
        }
        return rateLimitConfigMap.get(routeId);
    }
 
 
    @NotNull
    public Map<String, String> getHeaders(RateRedisRateLimiter.Config config, Long tokensLeft) {
        Map<String, String> headers = new HashMap();
        if (this.isIncludeHeaders()) {
            headers.put(this.remainingHeader, tokensLeft.toString());
            headers.put(this.replenishRateHeader, String.valueOf(config.getReplenishRate()));
            headers.put(this.burstCapacityHeader, String.valueOf(config.getBurstCapacity()));
        }
 
        return headers;
    }
 
    static List<String> getKeys(String id) {
        String prefix = "request_rate_limiter.{" + id;
        String tokenKey = prefix + "}.tokens";
        String timestampKey = prefix + "}.timestamp";
        return Arrays.asList(tokenKey, timestampKey);
    }

    @Validated
    public static class Config {
        @Min(1L)
        private int replenishRate;
        @Min(1L)
        private int burstCapacity = 1;

        private int timeType;

        private EnvironmentEnum environmentEnum;

        public Config() {
        }
 
        public int getReplenishRate() {
            return this.replenishRate;
        }
 
        public RateRedisRateLimiter.Config setReplenishRate(int replenishRate) {
            this.replenishRate = replenishRate;
            return this;
        }
 
        public int getBurstCapacity() {
            return this.burstCapacity;
        }
 
        public RateRedisRateLimiter.Config setBurstCapacity(int burstCapacity) {
            this.burstCapacity = burstCapacity;
            return this;
        }

        public int getTimeType() {
            return timeType;
        }

        public void setTimeType(int timeType) {
            this.timeType = timeType;
        }

        public EnvironmentEnum getEnvironmentEnum() {
            return environmentEnum;
        }

        public void setEnvironmentEnum(EnvironmentEnum environmentEnum) {
            this.environmentEnum = environmentEnum;
        }

        public String toString() {
            return "Config{replenishRate=" + this.replenishRate + ", burstCapacity=" + this.burstCapacity + '}';
        }
 
    }

    public static void main(String[] args) {
        System.out.println(Instant.now().getEpochSecond());
    }
}